import os
import numpy as np
from ultralytics import YOLO  # 导入 YOLOv8
from flask import Flask, request, send_file
from flask_cors import CORS
import cv2
import time

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})  # 允许所有跨域请求

# 设置基本路径和上传目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_DIR = os.path.join(BASE_DIR, 'img')
os.makedirs(UPLOAD_DIR, exist_ok=True)  # 创建上传目录

latest_result_path = None  # 全局变量，保存最新结果路径

@app.route('/request', methods=['GET', 'POST'])
def uploads():
    global latest_result_path

    # 获取前端上传的图片
    img = request.files.get('img')
    if not img:
        print("错误：没有上传图片")
        return "错误：没有上传图片", 400

    # 动态生成文件名
    timestamp = int(time.time() * 1000)
    name = f"img_{timestamp}.jpg"
    img_path = os.path.join(UPLOAD_DIR, name)

    # 保存上传的图片并验证
    try:
        img.save(img_path)
        if not os.path.exists(img_path):
            print(f"错误：图片保存失败 - {img_path}")
            return "错误：图片保存失败", 500
        print(f"图片已保存至：{img_path}")
    except Exception as e:
        print(f"保存图片时发生错误：{str(e)}")
        return f"错误：保存图片失败 - {str(e)}", 500

    # 使用 YOLO 模型处理图片
    try:
        results = model.predict(img_path, save=True, device=0)  # 处理图片并保存结果
        predict_dir = results[0].save_dir  # 获取 YOLO 输出目录
        print(f"预测结果保存目录：{predict_dir}")

        # 检查输出目录中是否有图片文件
        result_path = os.path.join(predict_dir, name)
        if not os.path.exists(result_path):
            print(f"警告：预期输出文件 {result_path} 不存在，检查实际输出")
            saved_files = os.listdir(predict_dir)
            for f in saved_files:
                if f.endswith('.jpg') or f.endswith('.png'):
                    result_path = os.path.join(predict_dir, f)
                    print(f"找到实际输出文件：{result_path}")
                    break
            else:
                print(f"错误：在 {predict_dir} 中未找到任何图片文件")
                return "错误：YOLO 处理未生成图片", 500

        latest_result_path = result_path
        print(f"处理结果路径：{latest_result_path}")
        return '成功'
    except Exception as e:
        print(f"YOLO 处理失败：{str(e)}")
        return f"错误：YOLO 处理失败 - {str(e)}", 500

@app.route('/get', methods=['GET', 'POST'])
def download():
    global latest_result_path
    if not latest_result_path:
        print("错误：latest_result_path 为空")
        return "错误：没有可用的结果", 404
    if not os.path.exists(latest_result_path):
        print(f"错误：文件不存在 - {latest_result_path}")
        return "错误：文件未找到", 404

    print(f"正在返回文件：{latest_result_path}")
    if latest_result_path.endswith('.jpg') or latest_result_path.endswith('.png'):
        return send_file(latest_result_path, mimetype='image/jpeg')
    else:
        print(f"错误：{latest_result_path} 不是图片文件")
        return "错误：结果不是图片", 400

if __name__ == '__main__':
    try:
        model = YOLO('./yolov8n.pt')  # 加载 YOLOv8 模型
        print("YOLO 模型加载成功")
        app.run(host='0.0.0.0', port=5000, debug=True)  # 启动 Flask 服务
    except Exception as e:
        print(f"启动失败：{str(e)}")